
import ode
import time
from time import sleep
from multiprocessing import Process, Pipe
from scipy import array
from math import sin, cos, pi

density = { "balsa" : 120,
			"brick" : 2000,
			"teflon" : 2200,
			"acrylic" : 1190,
			"copper" : 8900,
			"polystyrene" : 1050,
			"corc" : 250,
			"diamond" : 3300,
			"glass" : 2500,
			"gold" : 19300,
			"iron" : 7900,
			"lead" : 11300,
			"styrofoam" : 100 }


class PhysicalPart:
	def __init__(self, position, rotation, shape, 
				 material = None, size_xyz = None, 
				 radius = None, length = None):
		self.initial_position = position
		self.initial_rotation = rotation
		self.shape = shape
		self.material = material
		self.size_xyz = size_xyz
		self.radius = radius
		self.length = length
	
	def reset(self):
		self.position = self.initial_position
		self.rotation = self.initial_rotation


def setTransform(ode_body, physical_part):
	ode_body.setPosition(physical_part.position)
	"""
	omega = physical_part.rotation[0]
	phi = physical_part.rotation[1]
	theta = physical_part.rotation[2]
	
	omega = pi / 2.0
	
	rx = array([[1.0,			0.0,				0.0], 
				[0.0,			cos(omega),		-sin(omega)],
				[0.0,			sin(omega),		cos(omega)]])
	
	ry = array([[cos(phi),		0.0,			sin(phi)], 
				[0.0, 			1.0,			0.0],
				[-sin(phi),		0.0,			cos(phi)]])
	
	rz = array([[cos(theta), 	-sin(theta), 	0.0], 
				[sin(theta), 	cos(theta), 	0.0],
				[0.0, 			0.0, 			1.0]])
	
	rotation = [item for sublist in (rx * ry * rz).tolist() for item in sublist]
	rotation = (rx * ry * rz).tolist()
	r = rotation[0]
	r.extend(rotation[1])
	r.extend(rotation[2])
	ode_body.setRotation(r)
	print "---"
	print "o", r
	print ode_body.getRotation()
	print "---"
	
	rotation = [1.0, 0.0, 0.0, 0.0,0.0,0.0,0.0,0.0,0.0]
	print "o", rotation
	ode_body.setRotation(rotation)
	print ode_body.getRotation()
	
	quit()
	rotation = rotation[0].extend(rotation[1]).extend(rotation[2])
	print rotation
	rotation = [1.0, 0.0, 0.0, 0.0,0.0,0.0,0.0,0.0,0.0]
	print rotation
	ode_body.setRotation(rotation)
	print ode_body.getRotation()
	quit()
	
	
	print rx
	print ry
	print rz
	print rx * ry * rz
	print ode_body.isEnabled()
	rotation = [item for sublist in (rx * ry * rz).tolist() for item in sublist]
	print rotation
	#ode_body.setRotation(rotation)

	ode_body.setRotation([1.0, 1.0, 0.0, 1.0,0.0,0.0,0.0,0.0,-1.0])

	print ode_body.getRotation()
	quit()
	"""

# Collision callback
def near_callback(args, geom1, geom2):
    """Callback function for the collide() method.

    This function checks if the given geoms do collide and
    creates contact joints if they do.
    """
    # Check if the objects do collide
    contacts = ode.collide(geom1, geom2)

    # Create contact joints
    world,contactgroup = args
    for c in contacts:
        c.setBounce(0.2)
        c.setMu(500)
        j = ode.ContactJoint(world, contactgroup, c)
        j.attach(geom1.getBody(), geom2.getBody())


class Simulator:
	def __init__(self, physical_parts):		
		self.world = ode.World()
		self.world.setGravity((0.0, -9.81, 0.0))
		self.world.setERP(0.1)
		self.world.setCFM(1E-4)
		self.space = ode.Space()
		self.floor = ode.GeomPlane(self.space, (0,1,0), 0)
		self.fps = 50.0
		self.dt = 1.0/self.fps
		self.iterations_per_step = 5
		self.bodies = []

		for key in physical_parts:
			part = physical_parts[key]
			part.reset()
			if part.shape == "cylinder":				
				cylinder = self.create_cylinder(part)
				setTransform(cylinder, part)
				self.bodies.append(cylinder)
		
		

	def create_cylinder(self, part):
		assert(part.radius is not None)
		assert(part.length is not None)
		assert(part.material is not None)
		
		# Create body
		body = ode.Body(self.world)
		M = ode.Mass()
		M.setCylinder(density[part.material], 3, part.radius, part.length)
		body.setMass(M)
		
		# Create geometry for collision detection
		geometry = ode.GeomCylinder(self.space, part.radius, part.length)
		geometry.setBody(body)
		body.geom = geometry
		
		return body

	def create_box(self, part):
		assert(part.lx is not None)
		assert(part.ly is not None)
		assert(part.lz is not None)
		assert(part.material is not None)
		
		# Create body
		body = ode.Body(self.world)
		M = ode.Mass()
		M.setBox(density[part.material], part.lx, part.ly, part.lz)
		body.setMass(M)
		
		# Create geometry for collision detection
		geometry = ode.GeomBox(self.space, lengths = (part.lx, part.ly, part.lz))
		geometry.setBody(body)
		body.geom = geometry
		
		return body

	def start(self):
		self.parent_conn, child_conn = Pipe()
		p = Process(target=self._runner, args=(child_conn,))
		p.start()
	
	def stop(self):
		self.parent_conn.send(("stop",))
	
	def _runner(self, conn):
		lasttime = time.time()
		contactgroup = ode.JointGroup()
		
		while True:
			if conn.poll():
				data = conn.recv()
				if data[0] == "stop":
					print "stopping"
					break

			t = self.dt - (time.time() - lasttime)
			if (t > 0):
				sleep(t)
			lasttime = time.time()

			for step in range(self.iterations_per_step):
				# Detect collisions and create contact joints
				self.space.collide((self.world, contactgroup), near_callback)
			
				# Simulation step
				self.world.step(self.dt / self.iterations_per_step)

				for body in self.bodies:
					print body.getPosition()
				
				# Remove all contact joints
				contactgroup.empty()

			print "-----------------------"
			print "simulating"
	

if __name__ == "__main__":
	
	simulator = Simulator()
	simulator.simulate()
